from openai import OpenAI
from typing import Dict, Any, List
import json
import os

class CategoryItem:
    def __init__(self, id: str, name: str):
        self.id = id
        self.name = name

class ContentClassifier:
    def __init__(self, model_name: str = "qwen-2.5-72b"):
        """初始化配置"""
        self.client = OpenAI(
            api_key=os.getenv("OPENAI_API_KEY", "EMPTY"),
            base_url="http://172.16.8.36:9997/v1"
        )
        self.model_name = model_name

    def _generate_prompt(self, categories: List[CategoryItem]) -> str:
        """生成分类提示词"""
        category_list = [f"{cat.id}:{cat.name}" for cat in categories]
        
        prompt = f"""你是内容分类专家。职责：
分析内容，将其准确分类到预定义类别中，输出格式如下：
{{
    "分类结果": {{
        "id": "",       // 选中类别的ID
        "name": ""      // 选中类别的名称
    }},
    "分类依据": "",     // 详细说明选择该类别的原因
    "关键词": [],       // 3-5个支持分类判断的关键词或短语
    "置信度": 0.0      // 分类的置信程度（0-1之间）
}}

预定义类别：
{', '.join(category_list)}

注意事项：
1. 必须且只能从预定义类别中选择一个最贴切的类别
2. 详细说明选择该类别的依据
3. 给出分类的置信度评分（0-1之间）
4. 提取能够支持分类判断的关键词或短语
5. 确保分类结果中的id和name与预定义类别中的完全一致
"""
        return prompt

    def classify_content(self, content: str, categories: List[Dict[str, str]]) -> Dict[str, Any]:
        """对内容进行分类"""
        try:
            # 转换类别列表
            category_items = [CategoryItem(cat["id"], cat["name"]) for cat in categories]
            
            # 生成提示词
            prompt = self._generate_prompt(category_items)
            
            # 发送分类请求
            classify_response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": prompt},
                    {"role": "user", "content": content}
                ],
                temperature=0.1
            )
            
            classify_data = classify_response.choices[0].message.content
            if not classify_data:
                return None
                
            try:
                json_start = classify_data.find('{')
                json_end = classify_data.rfind('}') + 1
                if json_start >= 0 and json_end > json_start:
                    result = json.loads(classify_data[json_start:json_end])
                else:
                    return None
            except json.JSONDecodeError:
                print("分类结果不是有效的JSON格式")
                return None

            return result

        except Exception as e:
            print(f"处理内容时出错: {e}")
            return None

    def save_result(self, result: Dict[str, Any], output_dir: str, file_id: str) -> bool:
        """保存处理结果"""
        try:
            os.makedirs(output_dir, exist_ok=True)
            output_file = os.path.join(output_dir, f"{file_id}.json")
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"保存结果失败: {e}")
            return False

# 测试代码
if __name__ == "__main__":
    # 测试文章
    test_content = """
财政部决定开展国债做市支持操作
2024年11月18日 15:01 来源： 财政部网站 181人评论99+99+99+
摘要
【财政部决定开展国债做市支持操作】财政部办公厅发布通知，为支持国债做市，提高国债二级市场流动性，健全反映市场供求关系的国债收益率曲线，财政部决定开展国债做市支持操作。2024年11月22日起与同期国债合并上市交易。中标参与机构不迟于2024年11月20日，将随卖操作款缴入财政部指定资金账户，缴款日期以财政部指定账户收到款项为准。
　　在东方财富看资讯行情，选东方财富证券一站式开户交易>>

关于开展2024年11月份国债做市支持操作有关事宜的通知

财办库〔2024〕243号

　　国债做市支持参与机构，中国外汇交易中心暨全国银行间同业拆借中心，中央国债登记结算有限责任公司：

　　为支持国债做市，提高国债二级市场流动性，健全反映市场供求关系的国债收益率曲线，财政部决定开展国债做市支持操作（以下称本次操作）。现就有关事宜通知如下：

　　一、操作主要安排


　　二、投标限定


　　三、随卖券上市安排和操作款缴纳

　　2024年11月22日起与同期国债合并上市交易。中标参与机构不迟于2024年11月20日，将随卖操作款缴入财政部指定资金账户，缴款日期以财政部指定账户收到款项为准。

　　收款人名称：中华人民共和国财政部

　　开户银行：国家金库总库

　　汇入行行号：011100099992

　　2024年记账式附息（十六期）国债缴款账号：270-24116-10

　　2024年记账式附息（十九期）国债缴款账号：270-24119-10

　　四、其他

　　本次操作除遵循上述规定外，其他事宜按照《财政部中国人民银行关于印发<国债做市支持操作规则>的通知》（财库〔2016〕154号）等规定执行。

　　财政部办公厅

　　2024年11月12日
  """

    # 测试类别列表
    # test_categories = [
    #     {"id": "3aaad", "name": "财经-地产"},
    #     {"id": "dddde", "name": "财经-金融-股票"},
    #     {"id": "aadee", "name": "新闻-娱乐"}
    # ]
    test_categories = [
        {"id": "fin_stock", "name": "财经-股票"},
        {"id": "fin_property", "name": "财经-房地产"},
        {"id": "tech_ai", "name": "科技-人工智能"},
        {"id": "tech_digital", "name": "科技-数码产品"},
        {"id": "ent_movie", "name": "娱乐-电影"},
        {"id": "ent_celebrity", "name": "娱乐-明星"},
        {"id": "sports_basketball", "name": "体育-篮球"},
        {"id": "edu_policy", "name": "教育-政策"}
    ]

    # 初始化分类器
    classifier = ContentClassifier()
    
    # 处理内容
    result = classifier.classify_content(test_content, test_categories)
    if result:
        print("\n分类结果:")
        print(json.dumps(result, ensure_ascii=False, indent=2))
        
        # 保存结果
        classifier.save_result(result, "classified_contents", "test_content") 